import asyncio
import time
from .异常 import 超时


class UDP协议(asyncio.DatagramProtocol):
    def __init__(self):
        self.数据读取队列 = asyncio.Queue()
        self.传输: asyncio.DatagramTransport | None = None
        self.需要发送的数据 = []
        self.链接关闭 = False
        self.loop = asyncio.get_running_loop()

    def connection_made(self, 传输):
        self.传输 = 传输
        print("UDP服务器已启动")
        for 需要发送的数据 in self.需要发送的数据:
            self.write(*需要发送的数据)

    def datagram_received(self, 数据, 地址):
        self.数据读取队列.put_nowait((数据, 地址))

    def error_received(self, 异常):
        self.链接关闭 = True
        print(f"错误: {异常}")

    def connection_lost(self, 异常):
        self.链接关闭 = True
        print("服务器连接已关闭")

    async def read(self, 超时时间: float | None = None):
        if self.链接关闭:
            return
        if 超时时间 is not None:
            开始时间 = time.time()
            while time.time() - 开始时间 < 超时时间:
                if not self.数据读取队列.empty():
                    return self.数据读取队列.get_nowait()
                await asyncio.sleep(0.1)
            raise 超时()
        else:
            return await self.数据读取队列.get()

    def write(self, data, addr=None):
        if self.链接关闭:
            return False
        if self.传输 is None:
            self.需要发送的数据.append((data, addr))
            return True
        self.传输.sendto(data, addr)
        return True
